{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#default_exp torch_core"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "from fastai.imports import *\n",
    "from fastai.torch_imports import *"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from PIL import Image"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#hide\n",
    "from nbdev.showdoc import *"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "_all_ = ['progress_bar','master_bar']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "if torch.cuda.is_available():\n",
    "    if torch.cuda.current_device()==0:\n",
    "        def_gpu = int(os.environ.get('DEFAULT_GPU') or 0)\n",
    "        if torch.cuda.device_count()>=def_gpu: torch.cuda.set_device(def_gpu)\n",
    "    torch.backends.cudnn.benchmark = True"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Torch Core\n",
    "\n",
    "> Basic pytorch functions used in the fastai library"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Arrays and show"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "@delegates(plt.subplots, keep=True)\n",
    "def subplots(nrows=1, ncols=1, figsize=None, imsize=3, add_vert=0, **kwargs):\n",
    "    if figsize is None: figsize=(ncols*imsize, nrows*imsize+add_vert)\n",
    "    fig,ax = plt.subplots(nrows, ncols, figsize=figsize, **kwargs)\n",
    "    if nrows*ncols==1: ax = array([ax])\n",
    "    return fig,ax"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#hide\n",
    "_,axs = subplots()\n",
    "test_eq(axs.shape,[1])\n",
    "plt.close()\n",
    "_,axs = subplots(2,3)\n",
    "test_eq(axs.shape,[2,3])\n",
    "plt.close()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "def _fig_bounds(x):\n",
    "    r = x//32\n",
    "    return min(5, max(1,r))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "@delegates(plt.Axes.imshow, keep=True, but=['shape', 'imlim'])\n",
    "def show_image(im, ax=None, figsize=None, title=None, ctx=None, **kwargs):\n",
    "    \"Show a PIL or PyTorch image on `ax`.\"\n",
    "    # Handle pytorch axis order\n",
    "    if hasattrs(im, ('data','cpu','permute')):\n",
    "        im = im.data.cpu()\n",
    "        if im.shape[0]<5: im=im.permute(1,2,0)\n",
    "    elif not isinstance(im,np.ndarray): im=array(im)\n",
    "    # Handle 1-channel images\n",
    "    if im.shape[-1]==1: im=im[...,0]\n",
    "\n",
    "    ax = ifnone(ax,ctx)\n",
    "    if figsize is None: figsize = (_fig_bounds(im.shape[0]), _fig_bounds(im.shape[1]))\n",
    "    if ax is None: _,ax = plt.subplots(figsize=figsize)\n",
    "    ax.imshow(im, **kwargs)\n",
    "    if title is not None: ax.set_title(title)\n",
    "    ax.axis('off')\n",
    "    return ax"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "`show_image` can show PIL images..."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "im = Image.open(TEST_IMAGE_BW)\n",
    "ax = show_image(im, cmap=\"Greys\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "...and color images with standard `CHW` dim order..."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "im2 = np.array(Image.open(TEST_IMAGE))\n",
    "ax = show_image(im2, figsize=(2,2))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "...and color images with `HWC` dim order..."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "im3 = torch.as_tensor(im2).permute(2,0,1)\n",
    "ax = show_image(im3, figsize=(2,2))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "@delegates(show_image, keep=True)\n",
    "def show_titled_image(o, **kwargs):\n",
    "    \"Call `show_image` destructuring `o` to `(img,title)`\"\n",
    "    show_image(o[0], title=str(o[1]), **kwargs)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "show_titled_image((im3,'A puppy'), figsize=(2,2))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "@delegates(subplots)\n",
    "def show_images(ims, nrows=1, ncols=None, titles=None, **kwargs):\n",
    "    \"Show all images `ims` as subplots with `rows` using `titles`\"\n",
    "    if ncols is None: ncols = int(math.ceil(len(ims)/nrows))\n",
    "    if titles is None: titles = [None]*len(ims)\n",
    "    axs = subplots(nrows, ncols, **kwargs)[1].flat\n",
    "    for im,t,ax in zip(ims, titles, axs): show_image(im, ax=ax, title=t)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "show_images((im,im3), titles=('number','puppy'), imsize=2)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "`ArrayImage`, `ArrayImageBW` and `ArrayMask` are subclasses of `ndarray` that know how to show themselves."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "class ArrayBase(ndarray):\n",
    "    \"An `ndarray` that can modify casting behavior\"\n",
    "    @classmethod\n",
    "    def _before_cast(cls, x): return x if isinstance(x,ndarray) else array(x)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "class ArrayImageBase(ArrayBase):\n",
    "    \"Base class for arrays representing images\"\n",
    "    _show_args = {'cmap':'viridis'}\n",
    "    def show(self, ctx=None, **kwargs):\n",
    "        return show_image(self, ctx=ctx, **{**self._show_args, **kwargs})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "class ArrayImage(ArrayImageBase):\n",
    "    \"An array representing an image\"\n",
    "    pass"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "class ArrayImageBW(ArrayImage):\n",
    "    \"An array representing an image\"\n",
    "    _show_args = {'cmap':'Greys'}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "class ArrayMask(ArrayImageBase):\n",
    "    \"An array representing an image mask\"\n",
    "    _show_args = {'alpha':0.5, 'cmap':'tab20', 'interpolation':'nearest'}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "im = Image.open(TEST_IMAGE)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "im_t = cast(im, ArrayImage)\n",
    "test_eq(type(im_t), ArrayImage)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ax = im_t.show(figsize=(2,2))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "test_fig_exists(ax)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Basics"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "@patch\n",
    "def __array_eq__(self:Tensor,b):\n",
    "    return torch.equal(self,b) if self.dim() else self==b"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "def _array2tensor(x):\n",
    "    if x.dtype==np.uint16: x = x.astype(np.float32)\n",
    "    return torch.from_numpy(x)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "@use_kwargs_dict(dtype=None, device=None, requires_grad=False, pin_memory=False)\n",
    "def tensor(x, *rest, **kwargs):\n",
    "    \"Like `torch.as_tensor`, but handle lists too, and can pass multiple vector elements directly.\"\n",
    "    if len(rest): x = (x,)+rest\n",
    "    # There was a Pytorch bug in dataloader using num_workers>0. Haven't confirmed if fixed\n",
    "    # if isinstance(x, (tuple,list)) and len(x)==0: return tensor(0)\n",
    "    res = (x if isinstance(x, Tensor)\n",
    "           else torch.tensor(x, **kwargs) if isinstance(x, (tuple,list))\n",
    "           else _array2tensor(x) if isinstance(x, ndarray)\n",
    "           else as_tensor(x.values, **kwargs) if isinstance(x, (pd.Series, pd.DataFrame))\n",
    "           else as_tensor(x, **kwargs) if hasattr(x, '__array__') or is_iter(x)\n",
    "           else _array2tensor(array(x), **kwargs))\n",
    "    if res.dtype is torch.float64: return res.float()\n",
    "    return res"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "test_eq(tensor(torch.tensor([1,2,3])), torch.tensor([1,2,3]))\n",
    "test_eq(tensor(array([1,2,3])), torch.tensor([1,2,3]))\n",
    "test_eq(tensor(1,2,3), torch.tensor([1,2,3]))\n",
    "test_eq_type(tensor(1.0), torch.tensor(1.0))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "def set_seed(s, reproducible=False):\n",
    "    \"Set random seed for `random`, `torch`, and `numpy` (where available)\"\n",
    "    try: torch.manual_seed(s)\n",
    "    except NameError: pass\n",
    "    try: torch.cuda.manual_seed_all(s)\n",
    "    except NameError: pass\n",
    "    try: np.random.seed(s%(2**32-1))\n",
    "    except NameError: pass\n",
    "    random.seed(s)\n",
    "    if reproducible:\n",
    "        torch.backends.cudnn.deterministic = True\n",
    "        torch.backends.cudnn.benchmark = False"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "set_seed(2*33)\n",
    "a1 = np.random.random()\n",
    "a2 = torch.rand(())\n",
    "a3 = random.random()\n",
    "set_seed(2*33)\n",
    "b1 = np.random.random()\n",
    "b2 = torch.rand(())\n",
    "b3 = random.random()\n",
    "test_eq(a1,b1)\n",
    "test_eq(a2,b2)\n",
    "test_eq(a3,b3)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "def unsqueeze(x, dim=-1, n=1):\n",
    "    \"Same as `torch.unsqueeze` but can add `n` dims\"\n",
    "    for _ in range(n): x = x.unsqueeze(dim)\n",
    "    return x"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "t = tensor([1])\n",
    "t2 = unsqueeze(t, n=2)\n",
    "test_eq(t2,t[:,None,None])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "def unsqueeze_(x, dim=-1, n=1):\n",
    "    \"Same as `torch.unsqueeze_` but can add `n` dims\"\n",
    "    for _ in range(n): x.unsqueeze_(dim)\n",
    "    return x"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "t = tensor([1])\n",
    "unsqueeze_(t, n=2)\n",
    "test_eq(t, tensor([1]).view(1,1,1))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "def _fa_rebuild_tensor (cls, *args, **kwargs): return cls(torch._utils._rebuild_tensor_v2(*args, **kwargs))\n",
    "def _fa_rebuild_qtensor(cls, *args, **kwargs): return cls(torch._utils._rebuild_qtensor  (*args, **kwargs))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "def apply(func, x, *args, **kwargs):\n",
    "    \"Apply `func` recursively to `x`, passing on args\"\n",
    "    if is_listy(x): return type(x)([apply(func, o, *args, **kwargs) for o in x])\n",
    "    if isinstance(x,dict):  return {k: apply(func, v, *args, **kwargs) for k,v in x.items()}\n",
    "    res = func(x, *args, **kwargs)\n",
    "    return res if x is None else retain_type(res, x)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "def maybe_gather(x, axis=0):\n",
    "    \"Gather copies of `x` on `axis` (if training is distributed)\"\n",
    "    if num_distrib()<=1: return x\n",
    "    ndim = x.ndim\n",
    "    res = [x.new_zeros(*x.shape if ndim > 0 else (1,)) for _ in range(num_distrib())]\n",
    "    torch.distributed.all_gather(res, x.contiguous() if ndim > 0 else x[None])\n",
    "    return torch.cat(res, dim=axis) if ndim > 0 else torch.cat(res, dim=axis).mean()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "def to_detach(b, cpu=True, gather=True):\n",
    "    \"Recursively detach lists of tensors in `b `; put them on the CPU if `cpu=True`.\"\n",
    "    def _inner(x, cpu=True, gather=True):\n",
    "        if not isinstance(x,Tensor): return x\n",
    "        x = x.detach()\n",
    "        if gather: x = maybe_gather(x)\n",
    "        return x.cpu() if cpu else x\n",
    "    return apply(_inner, b, cpu=cpu, gather=gather)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "`gather` only applies during distributed training and the result tensor will be the one gathered across processes if `gather=True` (as a result, the batch size will be multiplied by the number of processes)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "def to_half(b):\n",
    "    \"Recursively map lists of tensors in `b ` to FP16.\"\n",
    "    return apply(lambda x: x.half() if torch.is_floating_point(x) else x, b)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "def to_float(b):\n",
    "    \"Recursively map lists of int tensors in `b ` to float.\"\n",
    "    return apply(lambda x: x.float() if torch.is_floating_point(x) else x, b)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "# None: True if available; True: error if not available; False: use CPU\n",
    "defaults.use_cuda = None"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "def default_device(use_cuda=-1):\n",
    "    \"Return or set default device; `use_cuda`: None - CUDA if available; True - error if not available; False - CPU\"\n",
    "    if use_cuda != -1: defaults.use_cuda=use_cuda\n",
    "    use = defaults.use_cuda or (torch.cuda.is_available() and defaults.use_cuda is None)\n",
    "    assert torch.cuda.is_available() or not use\n",
    "    return torch.device(torch.cuda.current_device()) if use else torch.device('cpu')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# cuda\n",
    "_td = torch.device(torch.cuda.current_device())\n",
    "test_eq(default_device(None), _td)\n",
    "test_eq(default_device(True), _td)\n",
    "test_eq(default_device(False), torch.device('cpu'))\n",
    "default_device(None);"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "def to_device(b, device=None):\n",
    "    \"Recursively put `b` on `device`.\"\n",
    "    if defaults.use_cuda==False: device='cpu'\n",
    "    elif device is None: device=default_device()\n",
    "    def _inner(o): return o.to(device, non_blocking=True) if isinstance(o,Tensor) else o.to_device(device) if hasattr(o, \"to_device\") else o\n",
    "    return apply(_inner, b)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "t = to_device((3,(tensor(3),tensor(2))))\n",
    "t1,(t2,t3) = t"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# cuda\n",
    "test_eq_type(t,(3,(tensor(3).cuda(),tensor(2).cuda())))\n",
    "test_eq(t2.type(), \"torch.cuda.LongTensor\")\n",
    "test_eq(t3.type(), \"torch.cuda.LongTensor\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "def to_cpu(b):\n",
    "    \"Recursively map lists of tensors in `b ` to the cpu.\"\n",
    "    return to_device(b,'cpu')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "t3 = to_cpu(t3)\n",
    "test_eq(t3.type(), \"torch.LongTensor\")\n",
    "test_eq(t3, 2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "def to_np(x):\n",
    "    \"Convert a tensor to a numpy array.\"\n",
    "    return apply(lambda o: o.data.cpu().numpy(), x)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "t3 = to_np(t3)\n",
    "test_eq(type(t3), np.ndarray)\n",
    "test_eq(t3, 2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "def to_concat(xs, dim=0):\n",
    "    \"Concat the element in `xs` (recursively if they are tuples/lists of tensors)\"\n",
    "    if not xs: return xs\n",
    "    if is_listy(xs[0]): return type(xs[0])([to_concat([x[i] for x in xs], dim=dim) for i in range_of(xs[0])])\n",
    "    if isinstance(xs[0],dict):  return {k: to_concat([x[k] for x in xs], dim=dim) for k in xs[0].keys()}\n",
    "    #We may receive xs that are not concatenable (inputs of a text classifier for instance),\n",
    "    #   in this case we return a big list\n",
    "    try:    return retain_type(torch.cat(xs, dim=dim), xs[0])\n",
    "    except: return sum([L(retain_type(o_.index_select(dim, tensor(i)).squeeze(dim), xs[0])\n",
    "                          for i in range_of(o_)) for o_ in xs], L())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "test_eq(to_concat([tensor([1,2]), tensor([3,4])]), tensor([1,2,3,4]))\n",
    "test_eq(to_concat([tensor([[1,2]]), tensor([[3,4]])], dim=1), tensor([[1,2,3,4]]))\n",
    "test_eq_type(to_concat([(tensor([1,2]), tensor([3,4])), (tensor([3,4]), tensor([5,6]))]), (tensor([1,2,3,4]), tensor([3,4,5,6])))\n",
    "test_eq_type(to_concat([[tensor([1,2]), tensor([3,4])], [tensor([3,4]), tensor([5,6])]]), [tensor([1,2,3,4]), tensor([3,4,5,6])])\n",
    "test_eq_type(to_concat([(tensor([1,2]),), (tensor([3,4]),)]), (tensor([1,2,3,4]),))\n",
    "\n",
    "test_eq(to_concat([tensor([[1,2]]), tensor([[3,4], [5,6]])], dim=1), [tensor([1]),tensor([3, 5]),tensor([4, 6])])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "test_eq(type(to_concat([dict(foo=tensor([1,2]), bar=tensor(3,4))])), dict)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Tensor subtypes"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "@patch\n",
    "def set_meta(self:Tensor, x, copy_meta=False):\n",
    "    \"Set all metadata in `__dict__`\"\n",
    "    if not hasattr(x,'__dict__'): return\n",
    "    d = x.__dict__\n",
    "    if copy_meta:\n",
    "        d = copy(d)\n",
    "        if '_meta' in d: d['_meta'] = copy(d['_meta'])\n",
    "    self.__dict__ = d"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "@patch\n",
    "def get_meta(self:Tensor, n, d=None):\n",
    "    \"Set `n` from `self._meta` if it exists and returns default `d` otherwise\"\n",
    "    return getattr(self, '_meta', {}).get(n, d)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "if not hasattr(torch,'as_subclass'):\n",
    "    setattr(torch, 'as_subclass', torch.Tensor.as_subclass)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "@patch\n",
    "def as_subclass(self:Tensor, typ):\n",
    "    \"Cast to `typ` and include `__dict__` and meta\"\n",
    "    return retain_meta(self, torch.as_subclass(self, typ))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "`Tensor.set_meta` and `Tensor.as_subclass` work together to maintain `_meta` after casting."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class _T(Tensor): pass\n",
    "t = tensor(1.).requires_grad_()\n",
    "t._meta = {'img_size': 1}\n",
    "t2 = t.as_subclass(_T)\n",
    "test_eq(t._meta, t2._meta)\n",
    "test_eq(t2.get_meta('img_size'), 1)\n",
    "assert(t2.requires_grad_)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "class TensorBase(Tensor):\n",
    "    def __new__(cls, x, **kwargs):\n",
    "        res = cast(tensor(x), cls)\n",
    "        if kwargs: res._meta = kwargs\n",
    "        return res\n",
    "\n",
    "    @classmethod\n",
    "    def _before_cast(cls, x): return tensor(x)\n",
    "\n",
    "    def __reduce_ex__(self,proto):\n",
    "        torch.utils.hooks.warn_if_has_hooks(self)\n",
    "        args = (type(self), self.storage(), self.storage_offset(), tuple(self.size()), self.stride())\n",
    "        if self.is_quantized: args = args + (self.q_scale(), self.q_zero_point())\n",
    "        f = _fa_rebuild_qtensor if self.is_quantized else  _fa_rebuild_tensor\n",
    "        return (f, args + (self.requires_grad, OrderedDict()))\n",
    "\n",
    "    def gi(self, i):\n",
    "        res = self[i]\n",
    "        return res.as_subclass(type(self)) if isinstance(res,Tensor) else res\n",
    "\n",
    "    def __repr__(self):\n",
    "        return re.sub('tensor', self.__class__.__name__, super().__repr__())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "def _patch_tb():\n",
    "    if getattr(TensorBase,'_patched',False): return\n",
    "    TensorBase._patched = True\n",
    "\n",
    "    def get_f(fn):\n",
    "        def _f(self, *args, **kwargs):\n",
    "            cls = self.__class__\n",
    "            res = getattr(super(TensorBase, self), fn)(*args, **kwargs)\n",
    "            return retain_type(res, self, copy_meta=True)\n",
    "        return _f\n",
    "\n",
    "    t = tensor([1])\n",
    "    skips = 'as_subclass imag real __getitem__ __class__ __deepcopy__ __delattr__ __dir__ __doc__ __getattribute__ __hash__ __init__ \\\n",
    "        __init_subclass__ __new__ __reduce__ __reduce_ex__ __repr__ __module__ __setstate__'.split()\n",
    "\n",
    "    for fn in dir(t):\n",
    "        if fn in skips: continue\n",
    "        f = getattr(t, fn)\n",
    "        if isinstance(f, (MethodWrapperType, BuiltinFunctionType, BuiltinMethodType, MethodType, FunctionType)):\n",
    "            setattr(TensorBase, fn, get_f(fn))\n",
    "\n",
    "_patch_tb()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#export \n",
    "class TensorCategory(TensorBase): pass"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#export \n",
    "class TensorMultiCategory(TensorCategory): pass"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class _T(TensorBase): pass"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "t = _T(range(5))\n",
    "test_eq(t[0], 0)\n",
    "test_eq_type(t.gi(0), _T(0))\n",
    "test_eq_type(t.gi(slice(2)), _T([0,1]))\n",
    "test_eq_type(t+1, _T(range(1,6)))\n",
    "test_eq(repr(t), '_T([0, 1, 2, 3, 4])')\n",
    "\n",
    "test_eq(type(pickle.loads(pickle.dumps(t))), _T)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "t = tensor([1,2,3])\n",
    "m = TensorBase([False,True,True])\n",
    "test_eq(t[m], tensor([2,3]))\n",
    "t = tensor([[1,2,3],[1,2,3]])\n",
    "m = cast(tensor([[False,True,True],\n",
    "                 [False,True,True]]), TensorBase)\n",
    "test_eq(t[m], tensor([2,3,2,3]))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "t = tensor([[1,2,3],[1,2,3]])\n",
    "t._meta = {'img_size': 1}\n",
    "t2 = cast(t, TensorBase)\n",
    "test_eq(t2._meta, t._meta)\n",
    "x = retain_type(tensor([4,5,6]), t2)\n",
    "test_eq(x._meta, t._meta)\n",
    "t3 = TensorBase([[1,2,3],[1,2,3]], img_size=1)\n",
    "test_eq(t3._meta, t._meta)\n",
    "t4 = t2+1\n",
    "t4._meta['img_size'] = 2\n",
    "test_eq(t2._meta, {'img_size': 1})\n",
    "test_eq(t4._meta, {'img_size': 2})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "class TensorImageBase(TensorBase):\n",
    "    _show_args = ArrayImageBase._show_args\n",
    "    def show(self, ctx=None, **kwargs):\n",
    "        return show_image(self, ctx=ctx, **{**self._show_args, **kwargs})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "class TensorImage(TensorImageBase): pass"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "class TensorImageBW(TensorImage): _show_args = ArrayImageBW._show_args"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "class TensorMask(TensorImageBase):\n",
    "    _show_args = ArrayMask._show_args\n",
    "\n",
    "    def show(self, ctx=None, **kwargs):\n",
    "        codes = self.get_meta('codes')\n",
    "        if codes is not None: kwargs = merge({'vmin': 1, 'vmax': len(codes)}, kwargs)\n",
    "        return super().show(ctx=ctx, **kwargs)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "im = Image.open(TEST_IMAGE)\n",
    "im_t = cast(array(im), TensorImage)\n",
    "test_eq(type(im_t), TensorImage)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "im_t2 = cast(tensor(1), TensorMask)\n",
    "test_eq(type(im_t2), TensorMask)\n",
    "test_eq(im_t2, tensor(1))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ax = im_t.show(figsize=(2,2))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "test_fig_exists(ax)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#hide (last test of to_concat)\n",
    "test_eq_type(to_concat([TensorImage([1,2]), TensorImage([3,4])]), TensorImage([1,2,3,4]))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "class TitledTensorScalar(TensorBase):\n",
    "    \"A tensor containing a scalar that has a `show` method\"\n",
    "    def show(self, **kwargs): show_title(self.item(), **kwargs)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## L -"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "@patch\n",
    "def tensored(self:L):\n",
    "    \"`mapped(tensor)`\"\n",
    "    return self.map(tensor)\n",
    "@patch\n",
    "def stack(self:L, dim=0):\n",
    "    \"Same as `torch.stack`\"\n",
    "    return torch.stack(list(self.tensored()), dim=dim)\n",
    "@patch\n",
    "def cat  (self:L, dim=0):\n",
    "    \"Same as `torch.cat`\"\n",
    "    return torch.cat  (list(self.tensored()), dim=dim)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "show_doc(L.tensored)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "There are shortcuts for `torch.stack` and `torch.cat` if your `L` contains tensors or something convertible. You can manually convert with `tensored`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "t = L(([1,2],[3,4]))\n",
    "test_eq(t.tensored(), [tensor(1,2),tensor(3,4)])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "show_doc(L.stack)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "test_eq(t.stack(), tensor([[1,2],[3,4]]))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "show_doc(L.cat)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "test_eq(t.cat(), tensor([1,2,3,4]))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Chunks"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "def concat(*ls):\n",
    "    \"Concatenate tensors, arrays, lists, or tuples\"\n",
    "    if not len(ls): return []\n",
    "    it = ls[0]\n",
    "    if isinstance(it,torch.Tensor): res = torch.cat(ls)\n",
    "    elif isinstance(it,ndarray): res = np.concatenate(ls)\n",
    "    else:\n",
    "        res = itertools.chain.from_iterable(map(L,ls))\n",
    "        if isinstance(it,(tuple,list)): res = type(it)(res)\n",
    "        else: res = L(res)\n",
    "    return retain_type(res, it)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "a,b,c = [1],[1,2],[1,1,2]\n",
    "test_eq(concat(a,b), c)\n",
    "test_eq_type(concat(tuple (a),tuple (b)), tuple (c))\n",
    "test_eq_type(concat(array (a),array (b)), array (c))\n",
    "test_eq_type(concat(tensor(a),tensor(b)), tensor(c))\n",
    "test_eq_type(concat(TensorBase(a),TensorBase(b)), TensorBase(c))\n",
    "test_eq_type(concat([1,1],1), [1,1,1])\n",
    "test_eq_type(concat(1,1,1), L(1,1,1))\n",
    "test_eq_type(concat(L(1,2),1), L(1,2,1))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "class Chunks:\n",
    "    \"Slice and int indexing into a list of lists\"\n",
    "    def __init__(self, chunks, lens=None):\n",
    "        self.chunks = chunks\n",
    "        self.lens = L(map(len,self.chunks) if lens is None else lens)\n",
    "        self.cumlens = np.cumsum(0+self.lens)\n",
    "        self.totlen = self.cumlens[-1]\n",
    "\n",
    "    def __getitem__(self,i):\n",
    "        if isinstance(i,slice): return retain_type(self.getslice(i), old=self.chunks[0])\n",
    "        di,idx = self.doc_idx(i)\n",
    "        return retain_type(self.chunks[di][idx], old=self.chunks[0])\n",
    "\n",
    "    def getslice(self, i):\n",
    "        st_d,st_i = self.doc_idx(ifnone(i.start,0))\n",
    "        en_d,en_i = self.doc_idx(ifnone(i.stop,self.totlen+1))\n",
    "        res = [self.chunks[st_d][st_i:(en_i if st_d==en_d else sys.maxsize)]]\n",
    "        for b in range(st_d+1,en_d): res.append(self.chunks[b])\n",
    "        if st_d!=en_d and en_d<len(self.chunks): res.append(self.chunks[en_d][:en_i])\n",
    "        return concat(*res)\n",
    "\n",
    "    def doc_idx(self, i):\n",
    "        if i<0: i=self.totlen+i # count from end\n",
    "        docidx = np.searchsorted(self.cumlens, i+1)-1\n",
    "        cl = self.cumlens[docidx]\n",
    "        return docidx,i-cl"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "docs = L(list(string.ascii_lowercase[a:b]) for a,b in ((0,3),(3,7),(7,8),(8,16),(16,24),(24,26)))\n",
    "\n",
    "b = Chunks(docs)\n",
    "test_eq([b[ o] for o in range(0,5)], ['a','b','c','d','e'])\n",
    "test_eq([b[-o] for o in range(1,6)], ['z','y','x','w','v'])\n",
    "test_eq(b[6:13], 'g,h,i,j,k,l,m'.split(','))\n",
    "test_eq(b[20:77], 'u,v,w,x,y,z'.split(','))\n",
    "test_eq(b[:5], 'a,b,c,d,e'.split(','))\n",
    "test_eq(b[:2], 'a,b'.split(','))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "t = torch.arange(26)\n",
    "docs = L(t[a:b] for a,b in ((0,3),(3,7),(7,8),(8,16),(16,24),(24,26)))\n",
    "b = Chunks(docs)\n",
    "test_eq([b[ o] for o in range(0,5)], range(0,5))\n",
    "test_eq([b[-o] for o in range(1,6)], [25,24,23,22,21])\n",
    "test_eq(b[6:13], torch.arange(6,13))\n",
    "test_eq(b[20:77], torch.arange(20,26))\n",
    "test_eq(b[:5], torch.arange(5))\n",
    "test_eq(b[:2], torch.arange(2))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "docs = L(TensorBase(t[a:b]) for a,b in ((0,3),(3,7),(7,8),(8,16),(16,24),(24,26)))\n",
    "b = Chunks(docs)\n",
    "test_eq_type(b[:2], TensorBase(range(2)))\n",
    "test_eq_type(b[:5], TensorBase(range(5)))\n",
    "test_eq_type(b[9:13], TensorBase(range(9,13)))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Simple types"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "def show_title(o, ax=None, ctx=None, label=None, color='black', **kwargs):\n",
    "    \"Set title of `ax` to `o`, or print `o` if `ax` is `None`\"\n",
    "    ax = ifnone(ax,ctx)\n",
    "    if ax is None: print(o)\n",
    "    elif hasattr(ax, 'set_title'):\n",
    "        t = ax.title.get_text()\n",
    "        if len(t) > 0: o = t+'\\n'+str(o)\n",
    "        ax.set_title(o, color=color)\n",
    "    elif isinstance(ax, pd.Series):\n",
    "        while label in ax: label += '_'\n",
    "        ax = ax.append(pd.Series({label: o}))\n",
    "    return ax"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "test_stdout(lambda: show_title(\"title\"), \"title\")\n",
    "# ensure that col names are unique when showing to a pandas series\n",
    "assert show_title(\"title\", ctx=pd.Series(dict(a=1)), label='a').equals(pd.Series(dict(a=1,a_='title')))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "class ShowTitle:\n",
    "    \"Base class that adds a simple `show`\"\n",
    "    _show_args = {'label': 'text'}\n",
    "    def show(self, ctx=None, **kwargs):\n",
    "        \"Show self\"\n",
    "        return show_title(str(self), ctx=ctx, **merge(self._show_args, kwargs))\n",
    "\n",
    "class TitledInt(Int, ShowTitle):\n",
    "    _show_args = {'label': 'text'}\n",
    "    def show(self, ctx=None, **kwargs):\n",
    "        \"Show self\"\n",
    "        return show_title(str(self), ctx=ctx, **merge(self._show_args, kwargs))\n",
    "\n",
    "class TitledFloat(Float, ShowTitle):\n",
    "    _show_args = {'label': 'text'}\n",
    "    def show(self, ctx=None, **kwargs):\n",
    "        \"Show self\"\n",
    "        return show_title(str(self), ctx=ctx, **merge(self._show_args, kwargs))\n",
    "\n",
    "class TitledStr(Str, ShowTitle):\n",
    "    _show_args = {'label': 'text'}\n",
    "    def show(self, ctx=None, **kwargs):\n",
    "        \"Show self\"\n",
    "        return show_title(str(self), ctx=ctx, **merge(self._show_args, kwargs))\n",
    "\n",
    "class TitledTuple(fastuple, ShowTitle):\n",
    "    _show_args = {'label': 'text'}\n",
    "    def show(self, ctx=None, **kwargs):\n",
    "        \"Show self\"\n",
    "        return show_title(str(self), ctx=ctx, **merge(self._show_args, kwargs))\n",
    "\n",
    "add_docs(TitledInt, \"An `int` with `show`\"); add_docs(TitledStr, \"An `str` with `show`\");\n",
    "add_docs(TitledFloat, \"A `float` with `show`\"); add_docs(TitledTuple, \"A `fastuple` with `show`\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "show_doc(TitledInt, title_level=3)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "show_doc(TitledStr, title_level=3)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "show_doc(TitledFloat, title_level=3)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "test_stdout(lambda: TitledStr('s').show(), 's')\n",
    "test_stdout(lambda: TitledInt(1).show(), '1')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "show_doc(TitledTuple, title_level=3)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#hide\n",
    "df = pd.DataFrame(index = range(1))\n",
    "row = df.iloc[0]\n",
    "x = TitledFloat(2.56)\n",
    "row = x.show(ctx=row, label='lbl')\n",
    "test_eq(float(row.lbl), 2.56)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "@patch\n",
    "def truncate(self:TitledStr, n):\n",
    "    \"Truncate self to `n`\"\n",
    "    words = self.split(' ')[:n]\n",
    "    return TitledStr(' '.join(words))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Other functions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "if not hasattr(pd.DataFrame,'_old_init'): pd.DataFrame._old_init = pd.DataFrame.__init__"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "@patch\n",
    "def __init__(self:pd.DataFrame, data=None, index=None, columns=None, dtype=None, copy=False):\n",
    "    if data is not None and isinstance(data, Tensor): data = to_np(data)\n",
    "    self._old_init(data, index=index, columns=columns, dtype=dtype, copy=copy)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "def get_empty_df(n):\n",
    "    \"Return `n` empty rows of a dataframe\"\n",
    "    df = pd.DataFrame(index = range(n))\n",
    "    return [df.iloc[i] for i in range(n)]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "def display_df(df):\n",
    "    \"Display `df` in a notebook or defaults to print\"\n",
    "    try: from IPython.display import display, HTML\n",
    "    except: return print(df)\n",
    "    display(HTML(df.to_html()))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "def get_first(c):\n",
    "    \"Get the first element of c, even if c is a dataframe\"\n",
    "    return getattr(c, 'iloc', c)[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "def one_param(m):\n",
    "    \"First parameter in `m`\"\n",
    "    return first(m.parameters())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "def item_find(x, idx=0):\n",
    "    \"Recursively takes the `idx`-th element of `x`\"\n",
    "    if is_listy(x): return item_find(x[idx])\n",
    "    if isinstance(x,dict):\n",
    "        key = list(x.keys())[idx] if isinstance(idx, int) else idx\n",
    "        return item_find(x[key])\n",
    "    return x"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "def find_device(b):\n",
    "    \"Recursively search the device of `b`.\"\n",
    "    return item_find(b).device"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "t2 = to_device(tensor(0))\n",
    "dev = default_device()\n",
    "test_eq(find_device(t2), dev)\n",
    "test_eq(find_device([t2,t2]), dev)\n",
    "test_eq(find_device({'a':t2,'b':t2}), dev)\n",
    "test_eq(find_device({'a':[[t2],[t2]],'b':t2}), dev)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "def find_bs(b):\n",
    "    \"Recursively search the batch size of `b`.\"\n",
    "    return item_find(b).shape[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "x = torch.randn(4,5)\n",
    "test_eq(find_bs(x), 4)\n",
    "test_eq(find_bs([x, x]), 4)\n",
    "test_eq(find_bs({'a':x,'b':x}), 4)\n",
    "test_eq(find_bs({'a':[[x],[x]],'b':x}), 4)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "def np_func(f):\n",
    "    \"Convert a function taking and returning numpy arrays to one taking and returning tensors\"\n",
    "    def _inner(*args, **kwargs):\n",
    "        nargs = [to_np(arg) if isinstance(arg,Tensor) else arg for arg in args]\n",
    "        return tensor(f(*nargs, **kwargs))\n",
    "    functools.update_wrapper(_inner, f)\n",
    "    return _inner"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This decorator is particularly useful for using numpy functions as fastai metrics, for instance:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from sklearn.metrics import f1_score"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "@np_func\n",
    "def f1(inp,targ): return f1_score(targ, inp)\n",
    "\n",
    "a1,a2 = array([0,1,1]),array([1,0,1])\n",
    "t = f1(tensor(a1),tensor(a2))\n",
    "test_eq(f1_score(a1,a2), t)\n",
    "assert isinstance(t,Tensor)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "class Module(nn.Module, metaclass=PrePostInitMeta):\n",
    "    \"Same as `nn.Module`, but no need for subclasses to call `super().__init__`\"\n",
    "    def __pre_init__(self, *args, **kwargs): super().__init__()\n",
    "    def __init__(self): pass"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "show_doc(Module, title_level=3)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class _T(Module):\n",
    "    def __init__(self): self.f = nn.Linear(1,1)\n",
    "    def forward(self,x): return self.f(x)\n",
    "\n",
    "t = _T()\n",
    "t(tensor([1.]))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# export\n",
    "from torch.nn.parallel import DistributedDataParallel"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# export\n",
    "def get_model(model):\n",
    "    \"Return the model maybe wrapped inside `model`.\"\n",
    "    return model.module if isinstance(model, (DistributedDataParallel, nn.DataParallel)) else model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# export\n",
    "def one_hot(x, c):\n",
    "    \"One-hot encode `x` with `c` classes.\"\n",
    "    res = torch.zeros(c, dtype=torch.uint8)\n",
    "    if isinstance(x, Tensor) and x.numel()>0: res[x] = 1.\n",
    "    else: res[list(L(x, use_list=None))] = 1.\n",
    "    return res"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "test_eq(one_hot([1,4], 5), tensor(0,1,0,0,1).byte())\n",
    "test_eq(one_hot(torch.tensor([]), 5), tensor(0,0,0,0,0).byte())\n",
    "test_eq(one_hot(2, 5), tensor(0,0,1,0,0).byte())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "def one_hot_decode(x, vocab=None):\n",
    "    return L(vocab[i] if vocab else i for i,x_ in enumerate(x) if x_==1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "test_eq(one_hot_decode(tensor(0,1,0,0,1)), [1,4])\n",
    "test_eq(one_hot_decode(tensor(0,0,0,0,0)), [   ])\n",
    "test_eq(one_hot_decode(tensor(0,0,1,0,0)), [2  ])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "def params(m):\n",
    "    \"Return all parameters of `m`\"\n",
    "    return [p for p in m.parameters()]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "def trainable_params(m):\n",
    "    \"Return all trainable parameters of `m`\"\n",
    "    return [p for p in m.parameters() if p.requires_grad]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "m = nn.Linear(4,5)\n",
    "test_eq(trainable_params(m), [m.weight, m.bias])\n",
    "m.weight.requires_grad_(False)\n",
    "test_eq(trainable_params(m), [m.bias])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "norm_types = (nn.BatchNorm1d, nn.BatchNorm2d, nn.BatchNorm3d, nn.InstanceNorm1d, nn.InstanceNorm2d, nn.InstanceNorm3d, nn.LayerNorm)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "def norm_bias_params(m, with_bias=True):\n",
    "    \"Return all bias and BatchNorm parameters\"\n",
    "    if isinstance(m, norm_types): return L(m.parameters())\n",
    "    res = L(m.children()).map(norm_bias_params, with_bias=with_bias).concat()\n",
    "    if with_bias and getattr(m, 'bias', None) is not None: res.append(m.bias)\n",
    "    return res"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for norm_func in [nn.BatchNorm1d, partial(nn.InstanceNorm1d, affine=True)]:\n",
    "    model = nn.Sequential(nn.Linear(10,20), norm_func(20), nn.Conv1d(3,4, 3))\n",
    "    test_eq(norm_bias_params(model), [model[0].bias, model[1].weight, model[1].bias, model[2].bias])\n",
    "    model = nn.ModuleList([nn.Linear(10,20, bias=False), nn.Sequential(norm_func(20), nn.Conv1d(3,4,3))])\n",
    "    test_eq(norm_bias_params(model), [model[1][0].weight, model[1][0].bias, model[1][1].bias])\n",
    "    model = nn.ModuleList([nn.Linear(10,20), nn.Sequential(norm_func(20), nn.Conv1d(3,4,3))])\n",
    "    test_eq(norm_bias_params(model, with_bias=False), [model[1][0].weight, model[1][0].bias])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "def batch_to_samples(b, max_n=10):\n",
    "    \"'Transposes' a batch to (at most `max_n`) samples\"\n",
    "    if isinstance(b, Tensor): return retain_types(list(b[:max_n]), [b])\n",
    "    else:\n",
    "        res = L(b).map(partial(batch_to_samples,max_n=max_n))\n",
    "        return retain_types(res.zip(), [b])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "t = tensor([1,2,3])\n",
    "test_eq(batch_to_samples([t,t+1], max_n=2), ([1,2],[2,3]))\n",
    "test_eq(batch_to_samples(tensor([1,2,3]), 10), [1, 2, 3])\n",
    "test_eq(batch_to_samples([tensor([1,2,3]), tensor([4,5,6])], 10), [(1, 4), (2, 5), (3, 6)])\n",
    "test_eq(batch_to_samples([tensor([1,2,3]), tensor([4,5,6])], 2), [(1, 4), (2, 5)])\n",
    "test_eq(batch_to_samples([tensor([1,2,3]), [tensor([4,5,6]),tensor([7,8,9])]], 10), \n",
    "        [(1, (4, 7)), (2, (5, 8)), (3, (6, 9))])\n",
    "test_eq(batch_to_samples([tensor([1,2,3]), [tensor([4,5,6]),tensor([7,8,9])]], 2), [(1, (4, 7)), (2, (5, 8))])\n",
    "\n",
    "t = fastuple(tensor([1,2,3]),TensorBase([2,3,4]))\n",
    "test_eq_type(batch_to_samples(t)[0][1], TensorBase(2))\n",
    "test_eq(batch_to_samples(t).map(type), [fastuple]*3)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "@patch\n",
    "def interp_1d(x:Tensor, xp, fp):\n",
    "    \"Same as `np.interp`\"\n",
    "    slopes = (fp[1:]-fp[:-1])/(xp[1:]-xp[:-1])\n",
    "    incx = fp[:-1] - (slopes*xp[:-1])\n",
    "    locs = (x[:,None]>=xp[None,:]).long().sum(1)-1\n",
    "    locs = locs.clamp(0,len(slopes)-1)\n",
    "    return slopes[locs]*x + incx[locs]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "brks = tensor(0,1,2,4,8,64).float()\n",
    "ys = tensor(range_of(brks)).float()\n",
    "ys /= ys[-1].item()\n",
    "pts = tensor(0.2,0.5,0.8,3,5,63)\n",
    "\n",
    "preds = pts.interp_1d(brks, ys)\n",
    "test_close(preds.numpy(), np.interp(pts.numpy(), brks.numpy(), ys.numpy()))\n",
    "\n",
    "plt.scatter(brks,ys)\n",
    "plt.scatter(pts,preds)\n",
    "plt.legend(['breaks','preds']);"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "@patch\n",
    "def pca(x:Tensor, k=2):\n",
    "    \"Compute PCA of `x` with `k` dimensions.\"\n",
    "    x = x-torch.mean(x,0)\n",
    "    U,S,V = torch.svd(x.t())\n",
    "    return torch.mm(x,U[:,:k])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# export\n",
    "def logit(x):\n",
    "    \"Logit of `x`, clamped to avoid inf.\"\n",
    "    x = x.clamp(1e-7, 1-1e-7)\n",
    "    return -(1/x-1).log()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "def num_distrib():\n",
    "    \"Return the number of processes in distributed training (if applicable).\"\n",
    "    return int(os.environ.get('WORLD_SIZE', 0))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "def rank_distrib():\n",
    "    \"Return the distributed rank of this process (if applicable).\"\n",
    "    return int(os.environ.get('RANK', 0))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "def distrib_barrier():\n",
    "    \"Place a synchronization barrier in distributed training so that ALL sub-processes in the pytorch process group must arrive here before proceeding.\"\n",
    "    if num_distrib() > 1 and torch.distributed.is_initialized(): torch.distributed.barrier()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "# Saving arrays requires pytables - optional dependency\n",
    "try: import tables\n",
    "except: pass"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "def _comp_filter(lib='lz4',lvl=3): return tables.Filters(complib=f'blosc:{lib}', complevel=lvl)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "@patch\n",
    "def save_array(p:Path, o, complib='lz4', lvl=3):\n",
    "    \"Save numpy array to a compressed `pytables` file, using compression level `lvl`\"\n",
    "    if isinstance(o,Tensor): o = to_np(o)\n",
    "    with tables.open_file(p, mode='w', filters=_comp_filter(lib=complib,lvl=lvl)) as f: f.create_carray('/', 'data', obj=o)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Compression lib can be any of: blosclz, lz4, lz4hc, snappy, zlib or zstd."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "@patch\n",
    "def load_array(p:Path):\n",
    "    \"Save numpy array to a `pytables` file\"\n",
    "    with tables.open_file(p, 'r') as f: return f.root.data.read()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "inspect.getdoc(load_array)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "str(inspect.signature(load_array))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "def base_doc(elt):\n",
    "    \"Print a base documentation of `elt`\"\n",
    "    name = getattr(elt, '__qualname__', getattr(elt, '__name__', ''))\n",
    "    print(f'{name}{inspect.signature(elt)}\\n{inspect.getdoc(elt)}\\n')\n",
    "    print('To get a prettier result with hyperlinks to source code and documentation, install nbdev: pip install nbdev')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "def doc(elt):\n",
    "    \"Try to use doc form nbdev and fall back to `base_doc`\"\n",
    "    try:\n",
    "        from nbdev.showdoc import doc\n",
    "        doc(elt)\n",
    "    except: base_doc(elt)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "def nested_reorder(t, idxs):\n",
    "    \"Reorder all tensors in `t` using `idxs`\"\n",
    "    if isinstance(t, (Tensor,L)): return t[idxs]\n",
    "    elif is_listy(t): return type(t)(nested_reorder(t_, idxs) for t_ in t)\n",
    "    if t is None: return t\n",
    "    raise TypeError(f\"Expected tensor, tuple, list or L but got {type(t)}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "x = tensor([0,1,2,3,4,5])\n",
    "idxs = tensor([2,5,1,0,3,4])\n",
    "test_eq_type(nested_reorder(([x], x), idxs), ([idxs], idxs))\n",
    "\n",
    "y = L(0,1,2,3,4,5)\n",
    "z = L(i.item() for i in idxs)\n",
    "test_eq_type(nested_reorder((y, x), idxs), (z,idxs))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Image helpers"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "def make_cross_image(bw=True):\n",
    "    \"Create a tensor containing a cross image, either `bw` (True) or color\"\n",
    "    if bw:\n",
    "        im = torch.zeros(5,5)\n",
    "        im[2,:] = 1.\n",
    "        im[:,2] = 1.\n",
    "    else:\n",
    "        im = torch.zeros(3,5,5)\n",
    "        im[0,2,:] = 1.\n",
    "        im[1,:,2] = 1.\n",
    "    return im"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.imshow(make_cross_image(), cmap=\"Greys\");"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.imshow(make_cross_image(False).permute(1,2,0));"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "def show_image_batch(b, show=show_titled_image, items=9, cols=3, figsize=None, **kwargs):\n",
    "    \"Display batch `b` in a grid of size `items` with `cols` width\"\n",
    "    if items<cols: cols=items\n",
    "    rows = (items+cols-1) // cols\n",
    "    if figsize is None: figsize = (cols*3, rows*3)\n",
    "    fig,axs = plt.subplots(rows, cols, figsize=figsize)\n",
    "    for *o,ax in zip(*to_cpu(b), axs.flatten()): show(o, ax=ax, **kwargs)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "show_image_batch(([Image.open(TEST_IMAGE_BW),Image.open(TEST_IMAGE)],['bw','color']), items=2)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Model init"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "def requires_grad(m):\n",
    "    \"Check if the first parameter of `m` requires grad or not\"\n",
    "    ps = list(m.parameters())\n",
    "    return ps[0].requires_grad if len(ps)>0 else False"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "tst = nn.Linear(4,5)\n",
    "assert requires_grad(tst)\n",
    "for p in tst.parameters(): p.requires_grad_(False)\n",
    "assert not requires_grad(tst)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "def init_default(m, func=nn.init.kaiming_normal_):\n",
    "    \"Initialize `m` weights with `func` and set `bias` to 0.\"\n",
    "    if func:\n",
    "        if hasattr(m, 'weight'): func(m.weight)\n",
    "        if hasattr(m, 'bias') and hasattr(m.bias, 'data'): m.bias.data.fill_(0.)\n",
    "    return m"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "tst = nn.Linear(4,5)\n",
    "tst.weight.data.uniform_(-1,1)\n",
    "tst.bias.data.uniform_(-1,1)\n",
    "tst = init_default(tst, func = lambda x: x.data.fill_(1.))\n",
    "test_eq(tst.weight, torch.ones(5,4))\n",
    "test_eq(tst.bias, torch.zeros(5))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "def cond_init(m, func):\n",
    "    \"Apply `init_default` to `m` unless it's a batchnorm module\"\n",
    "    if (not isinstance(m, norm_types)) and requires_grad(m): init_default(m, func)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "tst = nn.Linear(4,5)\n",
    "tst.weight.data.uniform_(-1,1)\n",
    "tst.bias.data.uniform_(-1,1)\n",
    "cond_init(tst, func = lambda x: x.data.fill_(1.))\n",
    "test_eq(tst.weight, torch.ones(5,4))\n",
    "test_eq(tst.bias, torch.zeros(5))\n",
    "\n",
    "tst = nn.BatchNorm2d(5)\n",
    "init = [tst.weight.clone(), tst.bias.clone()]\n",
    "cond_init(tst, func = lambda x: x.data.fill_(1.))\n",
    "test_eq(tst.weight, init[0])\n",
    "test_eq(tst.bias, init[1])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "def apply_leaf(m, f):\n",
    "    \"Apply `f` to children of `m`.\"\n",
    "    c = m.children()\n",
    "    if isinstance(m, nn.Module): f(m)\n",
    "    for l in c: apply_leaf(l,f)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "tst = nn.Sequential(nn.Linear(4,5), nn.Sequential(nn.Linear(4,5), nn.Linear(4,5)))\n",
    "apply_leaf(tst, partial(init_default, func=lambda x: x.data.fill_(1.)))\n",
    "for l in [tst[0], *tst[1]]: test_eq(l.weight, torch.ones(5,4))\n",
    "for l in [tst[0], *tst[1]]: test_eq(l.bias,   torch.zeros(5))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "def apply_init(m, func=nn.init.kaiming_normal_):\n",
    "    \"Initialize all non-batchnorm layers of `m` with `func`.\"\n",
    "    apply_leaf(m, partial(cond_init, func=func))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "tst = nn.Sequential(nn.Linear(4,5), nn.Sequential(nn.Linear(4,5), nn.BatchNorm1d(5)))\n",
    "init = [tst[1][1].weight.clone(), tst[1][1].bias.clone()]\n",
    "apply_init(tst, func=lambda x: x.data.fill_(1.))\n",
    "for l in [tst[0], tst[1][0]]: test_eq(l.weight, torch.ones(5,4))\n",
    "for l in [tst[0], tst[1][0]]: test_eq(l.bias,   torch.zeros(5))\n",
    "test_eq(tst[1][1].weight, init[0])\n",
    "test_eq(tst[1][1].bias,   init[1])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## autograd jit functions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "def script_use_ctx(f):\n",
    "    \"Decorator: create jit script and pass everything in `ctx.saved_variables to `f`, after `*args`\"\n",
    "    sf = torch.jit.script(f)\n",
    "    def _f(ctx, *args, **kwargs): return sf(*args, *ctx.saved_variables, **kwargs)\n",
    "    return update_wrapper(_f,f)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "def script_save_ctx(static, *argidx):\n",
    "    \"Decorator: create jit script and save args with indices `argidx` using `ctx.save_for_backward`\"\n",
    "    def _dec(f):\n",
    "        sf = torch.jit.script(f)\n",
    "        def _f(ctx, *args, **kwargs):\n",
    "            if argidx:\n",
    "                save = [args[o] for o in argidx]\n",
    "                ctx.save_for_backward(*save)\n",
    "            if not argidx: args = [ctx]+args\n",
    "            return sf(*args, **kwargs)\n",
    "        if static: _f = staticmethod(_f)\n",
    "        return update_wrapper(_f,f)\n",
    "    return _dec"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "def script_fwd(*argidx):\n",
    "    \"Decorator: create static jit script and save args with indices `argidx` using `ctx.save_for_backward`\"\n",
    "    return script_save_ctx(True, *argidx)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "def script_bwd(f):\n",
    "    \"Decorator: create static jit script and pass everything in `ctx.saved_variables to `f`, after `*args`\"\n",
    "    return staticmethod(script_use_ctx(f))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "def grad_module(cls):\n",
    "    \"Decorator: convert `cls` into an autograd function\"\n",
    "    class _c(nn.Module):\n",
    "        def forward(self, *args, **kwargs): return cls.apply(*args, **kwargs)\n",
    "    return _c"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Export -"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#hide\n",
    "from nbdev.export import notebook2script\n",
    "notebook2script()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "jupytext": {
   "split_at_heading": true
  },
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
